{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "adc8346b",
   "metadata": {},
   "source": [
    "# Advanced Visualizations for Fed-MVKM Analysis\n",
    "\n",
    "This notebook provides sophisticated visualization tools for analyzing Fed-MVKM algorithm's behavior using interactive and 3D visualizations.\n",
    "\n",
    "**Author:** Kristina P. Sinaga  \n",
    "**Date:** May 2024  \n",
    "**Version:** 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddf01d5c",
   "metadata": {},
   "source": [
    "## 1. Setup and Imports\n",
    "\n",
    "Let's import required libraries with a focus on advanced visualization packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "615c3b12",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import plotly.graph_objects as go\n",
    "import plotly.express as px\n",
    "from plotly.subplots import make_subplots\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import networkx as nx\n",
    "from scipy.stats import gaussian_kde\n",
    "from sklearn.manifold import TSNE, MDS\n",
    "from sklearn.decomposition import PCA\n",
    "import umap\n",
    "import pandas as pd\n",
    "from typing import List, Dict, Tuple, Optional\n",
    "import colorcet as cc  # For better color maps\n",
    "import plotly.io as pio\n",
    "\n",
    "# Set default theme for plotly\n",
    "pio.templates.default = \"plotly_white\"\n",
    "\n",
    "# Import Fed-MVKM related modules\n",
    "from mvkm_ed import FedMVKMED, FedMVKMEDConfig\n",
    "from mvkm_ed.utils import MVKMEDDataProcessor, MVKMEDMetrics\n",
    "\n",
    "# Configure matplotlib style\n",
    "plt.style.use('seaborn')\n",
    "sns.set_palette(cc.glasbey_dark)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f258803f",
   "metadata": {},
   "source": [
    "## 2. Enhanced Cluster Visualization Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dd13bab",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EnhancedClusterVisualizer:\n",
    "    \"\"\"Advanced cluster visualization with multiple dimensionality reduction methods.\"\"\"\n",
    "    \n",
    "    def __init__(self, views: List[np.ndarray], labels: np.ndarray, predictions: np.ndarray):\n",
    "        self.views = views\n",
    "        self.labels = labels\n",
    "        self.predictions = predictions\n",
    "        self.n_views = len(views)\n",
    "        self.n_clusters = len(np.unique(predictions))\n",
    "        \n",
    "    def plot_cluster_comparison(self, method='tsne', dims=2):\n",
    "        \"\"\"Compare true labels vs predictions using various embedding methods.\"\"\"\n",
    "        methods = {\n",
    "            'tsne': TSNE,\n",
    "            'umap': umap.UMAP,\n",
    "            'pca': PCA,\n",
    "            'mds': MDS\n",
    "        }\n",
    "        \n",
    "        reducer = methods[method.lower()](\n",
    "            n_components=dims,\n",
    "            random_state=42\n",
    "        )\n",
    "        \n",
    "        fig = make_subplots(\n",
    "            rows=self.n_views,\n",
    "            cols=2,\n",
    "            subplot_titles=[\n",
    "                f'View {i+1} - True Labels vs Predictions'\n",
    "                for i in range(self.n_views)\n",
    "            ],\n",
    "            specs=[[{'type': 'scene' if dims == 3 else None}]*2]*self.n_views\n",
    "        )\n",
    "        \n",
    "        for view_idx, view in enumerate(self.views):\n",
    "            # Compute embedding\n",
    "            embedding = reducer.fit_transform(view)\n",
    "            \n",
    "            # Create scatter plots\n",
    "            if dims == 3:\n",
    "                self._add_3d_scatter(fig, embedding, self.labels, \n",
    "                                   view_idx+1, 1, 'True Labels')\n",
    "                self._add_3d_scatter(fig, embedding, self.predictions, \n",
    "                                   view_idx+1, 2, 'Predictions')\n",
    "            else:\n",
    "                self._add_2d_scatter(fig, embedding, self.labels, \n",
    "                                   view_idx+1, 1, 'True Labels')\n",
    "                self._add_2d_scatter(fig, embedding, self.predictions, \n",
    "                                   view_idx+1, 2, 'Predictions')\n",
    "        \n",
    "        fig.update_layout(\n",
    "            height=400*self.n_views,\n",
    "            width=1200,\n",
    "            title_text=f\"Cluster Comparison using {method.upper()}\"\n",
    "        )\n",
    "        fig.show()\n",
    "        \n",
    "    def _add_2d_scatter(self, fig, embedding, labels, row, col, title):\n",
    "        \"\"\"Add 2D scatter plot to figure.\"\"\"\n",
    "        scatter = go.Scatter(\n",
    "            x=embedding[:, 0],\n",
    "            y=embedding[:, 1],\n",
    "            mode='markers',\n",
    "            marker=dict(\n",
    "                size=8,\n",
    "                color=labels,\n",
    "                colorscale='Viridis',\n",
    "                showscale=True,\n",
    "                colorbar=dict(title='Cluster')\n",
    "            ),\n",
    "            text=[f'Cluster {l}' for l in labels],\n",
    "            name=title\n",
    "        )\n",
    "        fig.add_trace(scatter, row=row, col=col)\n",
    "        fig.update_xaxes(title_text=\"Dimension 1\", row=row, col=col)\n",
    "        fig.update_yaxes(title_text=\"Dimension 2\", row=row, col=col)\n",
    "        \n",
    "    def _add_3d_scatter(self, fig, embedding, labels, row, col, title):\n",
    "        \"\"\"Add 3D scatter plot to figure.\"\"\"\n",
    "        scatter = go.Scatter3d(\n",
    "            x=embedding[:, 0],\n",
    "            y=embedding[:, 1],\n",
    "            z=embedding[:, 2],\n",
    "            mode='markers',\n",
    "            marker=dict(\n",
    "                size=4,\n",
    "                color=labels,\n",
    "                colorscale='Viridis',\n",
    "                showscale=True,\n",
    "                colorbar=dict(title='Cluster')\n",
    "            ),\n",
    "            text=[f'Cluster {l}' for l in labels],\n",
    "            name=title\n",
    "        )\n",
    "        fig.add_trace(scatter, row=row, col=col)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deb6cacb",
   "metadata": {},
   "source": [
    "## 3. Federation Process Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0fa0f3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FederationVisualizer:\n",
    "    \"\"\"Visualize the federation process and client interactions.\"\"\"\n",
    "    \n",
    "    def __init__(self, model: FedMVKMED, client_data: Dict):\n",
    "        self.model = model\n",
    "        self.client_data = client_data\n",
    "        self.history = model.history\n",
    "        \n",
    "    def plot_federation_topology(self):\n",
    "        \"\"\"Visualize federation topology and client relationships.\"\"\"\n",
    "        G = nx.Graph()\n",
    "        \n",
    "        # Add nodes for clients and server\n",
    "        G.add_node('Server', node_type='server')\n",
    "        for client_id in self.client_data.keys():\n",
    "            G.add_node(f'Client {client_id}', node_type='client')\n",
    "            G.add_edge('Server', f'Client {client_id}')\n",
    "        \n",
    "        # Create positions\n",
    "        pos = nx.spring_layout(G)\n",
    "        \n",
    "        # Create figure\n",
    "        plt.figure(figsize=(12, 8))\n",
    "        \n",
    "        # Draw nodes\n",
    "        nx.draw_networkx_nodes(G, pos, \n",
    "                             node_color=['red' if n == 'Server' else 'blue' \n",
    "                                       for n in G.nodes()],\n",
    "                             node_size=[2000 if n == 'Server' else 1000 \n",
    "                                      for n in G.nodes()],\n",
    "                             alpha=0.7)\n",
    "        \n",
    "        # Draw edges with arrow to show communication direction\n",
    "        nx.draw_networkx_edges(G, pos, edge_color='gray', \n",
    "                             arrows=True, arrowsize=20)\n",
    "        \n",
    "        # Add labels\n",
    "        nx.draw_networkx_labels(G, pos)\n",
    "        \n",
    "        plt.title(\"Federation Topology and Communication Flow\")\n",
    "        plt.axis('off')\n",
    "        plt.show()\n",
    "        \n",
    "    def animate_model_evolution(self):\n",
    "        \"\"\"Create animated visualization of model evolution.\"\"\"\n",
    "        if 'center_updates' not in self.history:\n",
    "            print(\"No center updates found in history\")\n",
    "            return\n",
    "            \n",
    "        # Create animation frames\n",
    "        frames = []\n",
    "        for i, centers in enumerate(self.history['center_updates']):\n",
    "            for view_idx, view_centers in enumerate(centers):\n",
    "                frame = go.Frame(\n",
    "                    data=[go.Scatter(\n",
    "                        x=view_centers[:, 0],\n",
    "                        y=view_centers[:, 1],\n",
    "                        mode='markers+text',\n",
    "                        marker=dict(size=15, symbol='star'),\n",
    "                        text=[f'C{j}' for j in range(len(view_centers))],\n",
    "                        name=f'View {view_idx+1}'\n",
    "                    )],\n",
    "                    name=f'Frame {i}'\n",
    "                )\n",
    "                frames.append(frame)\n",
    "        \n",
    "        # Create figure\n",
    "        fig = go.Figure(\n",
    "            frames=frames,\n",
    "            layout=go.Layout(\n",
    "                title=\"Model Evolution Animation\",\n",
    "                updatemenus=[{\n",
    "                    'type': 'buttons',\n",
    "                    'showactive': False,\n",
    "                    'buttons': [{\n",
    "                        'label': 'Play',\n",
    "                        'method': 'animate',\n",
    "                        'args': [None, {'frame': {'duration': 500}}]\n",
    "                    }]\n",
    "                }]\n",
    "            )\n",
    "        )\n",
    "        \n",
    "        fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84334ff1",
   "metadata": {},
   "source": [
    "## 4. Advanced Performance Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e23c8dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PerformanceAnalyzer:\n",
    "    \"\"\"Advanced performance analysis and visualization tools.\"\"\"\n",
    "    \n",
    "    def __init__(self, model: FedMVKMED, views: List[np.ndarray], \n",
    "                 true_labels: np.ndarray):\n",
    "        self.model = model\n",
    "        self.views = views\n",
    "        self.true_labels = true_labels\n",
    "        \n",
    "    def plot_cluster_stability(self):\n",
    "        \"\"\"Visualize cluster stability across iterations.\"\"\"\n",
    "        if not hasattr(self.model, 'history'):\n",
    "            print(\"No history available in model\")\n",
    "            return\n",
    "            \n",
    "        assignments_history = self.model.history.get('assignments', [])\n",
    "        if not assignments_history:\n",
    "            print(\"No assignment history available\")\n",
    "            return\n",
    "            \n",
    "        # Calculate stability metrics\n",
    "        stability_matrix = np.zeros((len(assignments_history), \n",
    "                                   len(assignments_history)))\n",
    "        \n",
    "        for i in range(len(assignments_history)):\n",
    "            for j in range(len(assignments_history)):\n",
    "                stability_matrix[i, j] = self._compute_assignment_similarity(\n",
    "                    assignments_history[i],\n",
    "                    assignments_history[j]\n",
    "                )\n",
    "        \n",
    "        # Create heatmap\n",
    "        fig = go.Figure(data=go.Heatmap(\n",
    "            z=stability_matrix,\n",
    "            colorscale='RdBu',\n",
    "            text=np.around(stability_matrix, decimals=2),\n",
    "            texttemplate=\"%{text}\",\n",
    "            textfont={\"size\": 10},\n",
    "            hoverongaps=False\n",
    "        ))\n",
    "        \n",
    "        fig.update_layout(\n",
    "            title=\"Cluster Assignment Stability Matrix\",\n",
    "            xaxis_title=\"Iteration j\",\n",
    "            yaxis_title=\"Iteration i\",\n",
    "            width=800,\n",
    "            height=800\n",
    "        )\n",
    "        \n",
    "        fig.show()\n",
    "        \n",
    "    def _compute_assignment_similarity(self, assign1, assign2):\n",
    "        \"\"\"Compute similarity between two cluster assignments.\"\"\"\n",
    "        from sklearn.metrics import adjusted_rand_score\n",
    "        return adjusted_rand_score(assign1, assign2)\n",
    "    \n",
    "    def plot_view_contribution(self):\n",
    "        \"\"\"Analyze and visualize the contribution of each view.\"\"\"\n",
    "        view_weights = np.array(self.model.history['view_weights'])\n",
    "        \n",
    "        fig = make_subplots(\n",
    "            rows=2, cols=2,\n",
    "            subplot_titles=(\n",
    "                'View Weight Evolution',\n",
    "                'Final View Importance',\n",
    "                'View Weight Stability',\n",
    "                'View Discriminative Power'\n",
    "            )\n",
    "        )\n",
    "        \n",
    "        # 1. View weight evolution\n",
    "        for i in range(view_weights.shape[1]):\n",
    "            fig.add_trace(\n",
    "                go.Scatter(\n",
    "                    y=view_weights[:, i],\n",
    "                    mode='lines',\n",
    "                    name=f'View {i+1}'\n",
    "                ),\n",
    "                row=1, col=1\n",
    "            )\n",
    "        \n",
    "        # 2. Final view importance (pie chart)\n",
    "        fig.add_trace(\n",
    "            go.Pie(\n",
    "                values=view_weights[-1],\n",
    "                labels=[f'View {i+1}' for i in range(view_weights.shape[1])],\n",
    "                hole=0.3\n",
    "            ),\n",
    "            row=1, col=2\n",
    "        )\n",
    "        \n",
    "        # 3. View weight stability (box plot)\n",
    "        fig.add_trace(\n",
    "            go.Box(\n",
    "                y=view_weights.flatten(),\n",
    "                x=[f'View {i+1}' for i in range(view_weights.shape[1])] * len(view_weights),\n",
    "                name='Weight Distribution'\n",
    "            ),\n",
    "            row=2, col=1\n",
    "        )\n",
    "        \n",
    "        # 4. View discriminative power\n",
    "        discriminative_power = []\n",
    "        for view in self.views:\n",
    "            from sklearn.metrics import silhouette_score\n",
    "            score = silhouette_score(view, self.true_labels)\n",
    "            discriminative_power.append(score)\n",
    "        \n",
    "        fig.add_trace(\n",
    "            go.Bar(\n",
    "                x=[f'View {i+1}' for i in range(len(self.views))],\n",
    "                y=discriminative_power,\n",
    "                name='Discriminative Power'\n",
    "            ),\n",
    "            row=2, col=2\n",
    "        )\n",
    "        \n",
    "        fig.update_layout(\n",
    "            height=800,\n",
    "            width=1000,\n",
    "            title_text=\"Comprehensive View Analysis\"\n",
    "        )\n",
    "        \n",
    "        fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0baaa2e9",
   "metadata": {},
   "source": [
    "## 5. Example Usage\n",
    "\n",
    "Let's demonstrate how to use these advanced visualization tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d0eaed0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and preprocess data (assuming we have the data from previous notebooks)\n",
    "processor = MVKMEDDataProcessor()\n",
    "views, labels = load_mat_data(data_dir)\n",
    "views = processor.preprocess_views(views)\n",
    "\n",
    "# Initialize visualizers\n",
    "cluster_viz = EnhancedClusterVisualizer(views, labels, model.predict(client_data))\n",
    "fed_viz = FederationVisualizer(model, client_data)\n",
    "perf_analyzer = PerformanceAnalyzer(model, views, labels)\n",
    "\n",
    "# Generate visualizations\n",
    "print(\"1. Enhanced Cluster Visualization:\")\n",
    "cluster_viz.plot_cluster_comparison(method='umap', dims=3)\n",
    "\n",
    "print(\"\\n2. Federation Process Visualization:\")\n",
    "fed_viz.plot_federation_topology()\n",
    "fed_viz.animate_model_evolution()\n",
    "\n",
    "print(\"\\n3. Advanced Performance Analysis:\")\n",
    "perf_analyzer.plot_cluster_stability()\n",
    "perf_analyzer.plot_view_contribution()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebb68d3a",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "This notebook provided advanced visualization tools for:\n",
    "1. Enhanced cluster analysis with multiple dimensionality reduction methods\n",
    "2. Interactive 3D visualizations of federation process\n",
    "3. Animated model evolution tracking\n",
    "4. Comprehensive performance analysis\n",
    "5. View contribution analysis\n",
    "\n",
    "These visualizations help in better understanding and analyzing the Fed-MVKM algorithm's behavior and performance."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
